{#-
    Copyright (c) 2020-2024 NVIDIA CORPORATION & AFFILIATES.
    Apache-2.0

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
-#}
#!/bin/bash

declare -r SCRIPT_NAME="$(basename "$0")"
declare -r SCRIPT_PATH="$(readlink -f "$0")"
declare -r SCRIPT_DIR="$(dirname "$SCRIPT_PATH")"

declare -r LOCKFILE="/tmp/mlxfwmanager-lock"

declare -r YES_PARAM="yes"
declare -r NO_PARAM="no"

declare -r VERBOSE_ERROR="1"
declare -r VERBOSE_WARNING="2"
declare -r VERBOSE_NOTICE="3"
declare -r VERBOSE_INFO="4"

declare -r VERBOSE_MAX="${VERBOSE_INFO}"
declare -r VERBOSE_MIN="${VERBOSE_ERROR}"

declare -r EXIT_SUCCESS="0"
declare -r EXIT_FAILURE="1"
declare -r FW_ALREADY_UPDATED_FAILURE="2"
declare -r FW_UPGRADE_IS_REQUIRED="10"

declare -r QUERY_XML="mlxfwmanager --query-format XML"
declare -r QUERY_CMD="mlxfwmanager --query"
declare -r LIST_CONTENT_CMD="mlxfwmanager --list-content"
declare -r BURN_CMD="mlxfwmanager -u -f -y"

declare -r QUERY_FILE="/tmp/mlxfwmanager-query.log"
declare -r LIST_CONTENT_FILE="/tmp/mlxfwmanager-list-content.log"

declare -r SPC1_ASIC="spc1"
declare -r SPC2_ASIC="spc2"
declare -r SPC3_ASIC="spc3"
declare -r SPC4_ASIC="spc4"
declare -r SPC5_ASIC="spc5"
declare -r BF3_NIC="bf3"
declare -r UNKN_ASIC="unknown"
declare -r UNKN_MST="unknown"

declare -rA FW_FILE_MAP=( \
    [$SPC1_ASIC]="fw-SPC.mfa" \
    [$SPC2_ASIC]="fw-SPC2.mfa" \
    [$SPC3_ASIC]="fw-SPC3.mfa" \
    [$SPC4_ASIC]="fw-SPC4.mfa" \
    [$SPC5_ASIC]="fw-SPC5.mfa" \
    [$BF3_NIC]="fw-BF3.mfa" \
)

DRY_RUN="${NO_PARAM}"
IMAGE_UPGRADE="${NO_PARAM}"
SYSLOG_LOGGER="${NO_PARAM}"
VERBOSE_LEVEL="${VERBOSE_MIN}"
MFT_DIAGNOSIS_FLAGS=""
RESET_CONFIG="${NO_PARAM}"

function PrintHelp() {
    echo
    echo "Usage: ./${SCRIPT_NAME} [OPTIONS]"
    echo
    echo "OPTIONS:"
    echo "  -u, --upgrade  Upgrade ASIC firmware using next boot image (useful after SONiC-To-SONiC update)"
    echo "  -s, --syslog   Use syslog logger (enabled when -u|--upgrade)"
    echo "  -v, --verbose  Verbose mode (enabled when -u|--upgrade)"
    echo "  -d, --dry-run  Compare the FW versions without installation. Return code "0" means the FW is up-to-date, return code "10" means an upgrade is required, otherwise an error is detected."
    echo "  -c, --clear-semaphore   Clear hw resources before updating firmware"
{% if sonic_asic_platform == "nvidia-bluefield" %}
    echo "  -r, --reset    Reset firmware configuration (NVIDIA BlueField platform only)"
{% endif %}
    echo "  -h, --help     Print help"
    echo
    echo "Examples:"
    echo "  ./${SCRIPT_NAME} --verbose"
    echo "  ./${SCRIPT_NAME} --upgrade"
{% if sonic_asic_platform == "nvidia-bluefield" %}
    echo "  ./${SCRIPT_NAME} --reset"
{% endif %}
    echo "  ./${SCRIPT_NAME} --help"
    echo
}

function ParseArguments() {
    while [ "$#" -ge "1" ]; do
        case "$1" in
            -u|--upgrade)
                IMAGE_UPGRADE="${YES_PARAM}"
                SYSLOG_LOGGER="${YES_PARAM}"
            ;;
            -v|--verbose)
                VERBOSE_LEVEL="${VERBOSE_MAX}"
                MFT_DIAGNOSIS_FLAGS="FLASH_ACCESS_DEBUG=1 FW_COMPS_DEBUG=1"
            ;;
            -s|--syslog)
                SYSLOG_LOGGER="${YES_PARAM}"
            ;;
            -d|--dry-run)
                DRY_RUN="${YES_PARAM}"
            ;;
            -c|--clear-semaphore)
                CLEAR_SEMAPHORE="${YES_PARAM}"
            ;;
{% if sonic_asic_platform == "nvidia-bluefield" %}
            -r|--reset)
                RESET_CONFIG="${YES_PARAM}"
            ;;
{% endif %}
            -h|--help)
                PrintHelp
                exit "${EXIT_SUCCESS}"
            ;;
        esac
        shift
    done
}

function LogError() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_ERROR}" ]]; then
        echo "ERROR: $*"
        logger -p "ERROR" -t "${SCRIPT_NAME}" "$*"
    fi

    if [[ "${SYSLOG_LOGGER}" = "${YES_PARAM}" ]]; then
        logger -p "ERROR" -t "${SCRIPT_NAME}" "$*"
    fi
}

function LogWarning() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_WARNING}" ]]; then
        echo "WARNING: $*"
    fi

    if [[ "${SYSLOG_LOGGER}" = "${YES_PARAM}" ]]; then
        logger -p "WARNING" -t "${SCRIPT_NAME}" "$*"
    fi
}

function LogNotice() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_NOTICE}" ]]; then
        echo "NOTICE: $*"
    fi

    if [[ "${SYSLOG_LOGGER}" = "${YES_PARAM}" ]]; then
        logger -p "NOTICE" -t "${SCRIPT_NAME}" "$*"
    fi
}

function LogInfo() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_INFO}" ]]; then
        echo "INFO: $*"
    fi

    if [[ "${SYSLOG_LOGGER}" = "${YES_PARAM}" ]]; then
        logger -p "INFO" -t "${SCRIPT_NAME}" "$*"
    fi
}

function ExitFailure() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_ERROR}" ]]; then
        echo
        LogError "$@"
        echo
    fi

    exit "${EXIT_FAILURE}"
}

function ExitSuccess() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_INFO}" ]]; then
        echo
        LogInfo "$@"
        echo
    fi

    exit "${EXIT_SUCCESS}"
}

function LockStateChange() {
    LogInfo "Locking ${LOCKFILE} from ${SCRIPT_NAME}"

    exec {LOCKFD}>${LOCKFILE}
    /usr/bin/flock -x ${LOCKFD}

    LogInfo "Locked ${LOCKFILE} (${LOCKFD}) from ${SCRIPT_NAME}"
}

function UnlockStateChange() {
    LogInfo "Unlocking ${LOCKFILE} (${LOCKFD}) from ${SCRIPT_NAME}"
    /usr/bin/flock -u ${LOCKFD}
}

function GetMstDeviceType() {
    local -r asic_type=$(/usr/bin/asic_detect/asic_detect.sh)

    case $asic_type in
    ${SPC1_ASIC}|${SPC2_ASIC}|${SPC3_ASIC}|${SPC4_ASIC}|${SPC5_ASIC})
        echo "Spectrum"
        return ${EXIT_SUCCESS}
        ;;
    ${BF3_NIC})
        echo "BlueField3"
        return ${EXIT_SUCCESS}
        ;;
    *)
        echo "Unknown"
        return ${EXIT_FAILURE}
        ;;
    esac
}

function WaitForDevice() {
    local -i QUERY_RETRY_COUNT_MAX="10"
    local -i QUERY_RETRY_COUNT="0"
    local -r DEVICE_TYPE=$(GetMstDeviceType)
    local SPC_MST_DEV
    local QUERY_RC=""

    while : ; do
        SPC_MST_DEV=$(GetSPCMstDevice)
        ${QUERY_XML} -d ${SPC_MST_DEV} -o ${QUERY_FILE}
        QUERY_RC="$?"
        [[ ("${QUERY_RETRY_COUNT}" -lt "${QUERY_RETRY_COUNT_MAX}") && ("${QUERY_RC}" != "${EXIT_SUCCESS}") ]] || break
        sleep 1s
        ((QUERY_RETRY_COUNT++))
        LogInfo "Retrying MST device query ${QUERY_RETRY_COUNT}"
    done

    if [[ "${QUERY_RC}" != "${EXIT_SUCCESS}" ]]; then
        # Couldn't Detect the Spectrum ASIC. Exit failure and print the detailed information
        output=$(${QUERY_CMD})
        failure_msg="${output#*Fail : }"
        ExitFailure "FW Query command: ${QUERY_CMD} failed to detect spectrum device with error: ${failure_msg}"
    fi

    LogInfo "${DEVICE_TYPE} ASIC successfully detected at ${SPC_MST_DEV}"
}

function GetSPCMstDevice() {
    local _DEVICE_TYPE=$(GetMstDeviceType)
    local _MST_DEVICE=$(${QUERY_XML} | xmlstarlet sel -t -m "//Device[contains(@type,'${_DEVICE_TYPE}')]" -v @pciName | head -n 1)

    if [[ ! -c "${_MST_DEVICE}" ]]; then
        echo "${UNKN_MST}"
    else
        echo "${_MST_DEVICE}"
    fi

    exit "${EXIT_SUCCESS}"
}

function GetXPathXML() {
    local xpath=$1
    local xml_file=$2

    val=$(xmlstarlet sel -t -v "${xpath}" ${xml_file})
    ERROR_CODE="$?"
    if [[ "${ERROR_CODE}" != "${EXIT_SUCCESS}" ]]; then
        ExitFailure "XML Fetch failed for path: ${xpath}, file: $(cat ${xml_file})"
    fi

    echo ${val}
}

function RunCmd() {
    local ERROR_CODE="${EXIT_SUCCESS}"

    if [[ "${VERBOSE_LEVEL}" -eq "${VERBOSE_MAX}" ]]; then
        eval "$@"
    else
        eval "$@" &>/dev/null
    fi

    ERROR_CODE="$?"
    if [[ "${ERROR_CODE}" != "${EXIT_SUCCESS}" ]]; then
        ExitFailure "command failed: $@"
    fi
}

{% if sonic_asic_platform == "mellanox" %}

function RunFwUpdateCmd() {
    local ERROR_CODE="${EXIT_SUCCESS}"
    local COMMAND="${MFT_DIAGNOSIS_FLAGS} ${BURN_CMD} $@"

    if [[ "${VERBOSE_LEVEL}" -eq "${VERBOSE_MAX}" ]]; then
        output=$(eval "${COMMAND}")
    else
        output=$(eval "${COMMAND}") >/dev/null 2>&1
    fi

    ERROR_CODE="$?"

    if [[ "${ERROR_CODE}" == "${FW_ALREADY_UPDATED_FAILURE}" ]]; then
        LogInfo "FW reactivation is required. Reactivating and updating FW ..."
        local -r _MST_DEVICE="$(GetSPCMstDevice)"
        local -r _CMD="flint -d ${_MST_DEVICE} ir"
        output=$(eval "${_CMD}")

        if [[ "${VERBOSE_LEVEL}" -eq "${VERBOSE_MAX}" ]]; then
            output=$(eval "${COMMAND}")
        else
            output=$(eval "${COMMAND}") >/dev/null 2>&1
        fi
    fi

    ERROR_CODE="$?"

    if [[ "${ERROR_CODE}" != "${EXIT_SUCCESS}" ]]; then
        echo "${output}"
        failure_msg="${output#*Fail : }"
        ExitFailure "FW Update command: ${COMMAND} failed with error: ${failure_msg}"
    fi
}

function GetAvailableFwVersion() {
    local -r _FW_FILE="$1"
    local -r _MST_DEVICE="$2"
    local -r _PSID="$3"

    RunCmd "${LIST_CONTENT_CMD} -i ${_FW_FILE} -d ${_MST_DEVICE} -o ${LIST_CONTENT_FILE}"

    local -r _FW_AVAILABLE_INFO="$(grep ${_PSID} ${LIST_CONTENT_FILE})"
    local -r _FW_AVAILABLE="$(echo ${_FW_AVAILABLE_INFO} | awk '{print $4}')"

    echo ${_FW_AVAILABLE}
}

{% elif sonic_asic_platform == "nvidia-bluefield" %}

function RunFwUpdateCmd() {
    local ERROR_CODE="${EXIT_SUCCESS}"
    local -r _MST_DEVICE="$(GetSPCMstDevice)"

    # Reactivate FW prior to burning...
    eval "flint -d ${_MST_DEVICE} ir" 2>&1 >/dev/null

    RunCmd "${MFT_DIAGNOSIS_FLAGS} flint $@ burn"
    RunCmd "mlxconfig -d ${_MST_DEVICE} -y r"
}

function GetAvailableFwVersion() {
    local -r _FW_FILE="$1"
    local -r _MST_DEVICE="$2"
    local -r _PSID="$3"

    RunCmd "flint -i ${_FW_FILE} --psid ${_PSID} query 2>&1 > ${LIST_CONTENT_FILE}"

    local -r _FW_AVAILABLE_INFO="$(grep 'FW Version:' ${LIST_CONTENT_FILE})"
    local -r _FW_AVAILABLE="$(echo ${_FW_AVAILABLE_INFO} | awk '{print $3}')"

    echo ${_FW_AVAILABLE}
}

{% endif %}

function UpgradeFW() {
    local -r _FW_BIN_PATH="$1"

    local -r _ASIC_TYPE="$(/usr/bin/asic_detect/asic_detect.sh)"
    if [[ "${_ASIC_TYPE}" = "${UNKN_ASIC}" ]]; then
        ExitFailure "failed to detect ASIC type"
    fi

    if [ ! -z "${_FW_BIN_PATH}" ]; then
        local -r _FW_FILE="${_FW_BIN_PATH}/${FW_FILE_MAP[$_ASIC_TYPE]}"
    else
        local -r _FW_FILE="/etc/mlnx/${FW_FILE_MAP[$_ASIC_TYPE]}"
    fi

    if [ ! -f "${_FW_FILE}" ]; then
        ExitFailure "no such file: ${_FW_FILE}"
    fi

    local -r _MST_DEVICE=$(GetSPCMstDevice)
    RunCmd "${QUERY_XML} -d ${_MST_DEVICE} -o ${QUERY_FILE}"
    local -r _FW_CURRENT=$(GetXPathXML "//Device/Versions/FW/@current" ${QUERY_FILE})
    local -r _PSID=$(GetXPathXML "//Device/@psid" ${QUERY_FILE})
    local -r _FW_AVAILABLE="$(GetAvailableFwVersion ${_FW_FILE} ${_MST_DEVICE} ${_PSID})"

    if [[ -z "${_FW_CURRENT}" ]]; then
        ExitFailure "could not retreive current FW version"
    fi

    if [[ -z "${_FW_AVAILABLE}" ]]; then
        ExitFailure "could not retreive available FW version"
    fi

    if [[ "${_FW_CURRENT}" == "${_FW_AVAILABLE}" ]]; then
        ExitSuccess "firmware is up to date"
    else
        if [[ "${DRY_RUN}" == "${YES_PARAM}" ]]; then
            LogNotice "firmware upgrade is required"
            exit ${FW_UPGRADE_IS_REQUIRED}
        fi

        LogNotice "firmware upgrade is required. Installing compatible version..."

        if [[ "${_MST_DEVICE}" = "${UNKN_MST}" ]]; then
            LogWarning "could not find fastest mst device, using default device"
            RunFwUpdateCmd "-i ${_FW_FILE}"
        else
            RunFwUpdateCmd "-d ${_MST_DEVICE} -i ${_FW_FILE}"
        fi
    fi
}

function UpgradeFWFromImage() {
    local -r _NEXT_SONIC_IMAGE="$(sonic-installer list | grep "Next: " | cut -f2 -d' ')"
    local -r _CURRENT_SONIC_IMAGE="$(sonic-installer list | grep "Current: " | cut -f2 -d' ')"

    if [[ "${_CURRENT_SONIC_IMAGE}" == "${_NEXT_SONIC_IMAGE}" ]]; then
        ExitSuccess "firmware is up to date"
    fi

    # /host/image-<version>/platform/fw/asic is now the new location for FW binaries.
    # Prefere this path and if it does not exist use squashfs as a fallback.
    local -r _PLATFORM_FW_BIN_PATH="/host/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}/platform/fw/asic/"

    if [[ -d "${_PLATFORM_FW_BIN_PATH}" ]]; then
        LogInfo "Using FW binaries from ${_PLATFORM_FW_BIN_PATH}"

        UpgradeFW "${_PLATFORM_FW_BIN_PATH}"
    else
        local -r _FS_PATH="/host/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}/fs.squashfs"
        local -r _FS_MOUNTPOINT="/tmp/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}-fs"
        local -r _FW_BIN_PATH="${_FS_MOUNTPOINT}/etc/mlnx/"

        LogInfo "Using FW binaries from ${_FW_BIN_PATH}"

        mkdir -p "${_FS_MOUNTPOINT}"
        mount -t squashfs "${_FS_PATH}" "${_FS_MOUNTPOINT}"

        UpgradeFW "${_FW_BIN_PATH}"

        umount -rf "${_FS_MOUNTPOINT}"
        rm -rf "${_FS_MOUNTPOINT}"
    fi
}

function ExitIfQEMU() {
    if [ -n "$(lspci -vvv | grep SimX)" ]; then
        ExitSuccess "No FW upgrade for SimX platform"
    fi
}

function Cleanup() {
    if [[ -n "${LOCKFD}" ]]; then
        UnlockStateChange
    fi
}

function ClearSemaphore() {
    if [[ "${CLEAR_SEMAPHORE}" == "${YES_PARAM}" ]]; then
        local -r _MST_DEVICE="$(GetSPCMstDevice)"
        if [[ "${_MST_DEVICE}" != "${UNKN_MST}" ]]; then
            /usr/bin/flint -d $_MST_DEVICE --clear_semaphore
        fi
    fi
}

function ResetFirmwareConfig() {
    local -r _MST_DEVICE="$(GetSPCMstDevice)"

    if [[ "${_MST_DEVICE}" = "${UNKN_MST}" ]]; then
        ExitFailure "Could not find MST device for firmware reset"
    fi

    LogInfo "Resetting firmware configuration for device ${_MST_DEVICE}"
    RunCmd "mlxconfig -d ${_MST_DEVICE} -y r"
    LogInfo "Firmware configuration reset completed successfully"
}

trap Cleanup EXIT

ParseArguments "$@"

ExitIfQEMU

LockStateChange

WaitForDevice

ClearSemaphore

if [[ "${RESET_CONFIG}" = "${YES_PARAM}" ]]; then
    ResetFirmwareConfig
    ExitSuccess "firmware configuration reset completed"
fi

if [ "${IMAGE_UPGRADE}" != "${YES_PARAM}" ]; then
    UpgradeFW
else
    UpgradeFWFromImage
fi

ExitSuccess "firmware upgrade is completed"
